import os
import random
import copy

import numpy as np
import cv2
import rclpy
from rclpy.node import Node
from loguru import logger
import torch
# from numpy import random
import time
import signal
from ais_bench.infer.interface import InferSession
from torchvision import transforms
import aclruntime
import numpy as np
from cv_bridge import CvBridge
from std_msgs.msg import String
from sensor_msgs.msg import Image, CompressedImage
from vision_msgs.msg import Detection2DArray, ObjectHypothesisWithPose, Detection2D

from road_class import RoadClass


class RoadClassification(Node):  # for inference
    def __init__(self, name):
        super().__init__(name)
        self.get_logger().info("大家好, 我是%s节点！" % name)
        self.bridge = CvBridge()
        self.bridge2 = CvBridge()
        
        self.get_logger().info("初始化road_classification_node成功！")

        # 创建订阅者
        self.subpath = self.create_subscription(
            Image, "Image_camera", self.recv_image_callback, 1)
        self.road_result_pub = self.create_publisher(
            String, "cmd_switch_gait", 1)

        self.string_msg = String()
        self.pred_class_msg_queue = []
        self.cmd = String()
        
        self.rc = RoadClass()
    

    def recv_image_callback(self, msg: Image):
        image = self.bridge.imgmsg_to_cv2(msg)
        self.cmd.data = self.rc.infer(image)
        logger.info(f"==========road_classification:  {self.cmd.data}")
        self.road_result_pub.publish(self.cmd)
       
            
    def run(self, image_file):
        img = self.preprocess(image_file)
        pred = self.om_sess.infer([img],custom_sizes=1000)[0]
        pred_class = self.postpreocess(pred)
        return pred_class


def main(args=None):
    rclpy.init(args=args)
    road_classification_node = RoadClassification("road_classification_node")
    rclpy.spin(road_classification_node)
    rclpy.shutdown()


if __name__ == '__main__':
    main()
